import { create } from "zustand";
import type { WSClient } from "@/utils/ws";

interface WSStore {
  wsMap: Map<string, WSClient>;
  setWS: (sessionId: string, ws: WSClient) => void;
  closeWS: (sessionId: string) => void;
}

export const useWSStore = create<WSStore>((set, get) => ({
  wsMap: new Map(),
  setWS: (sessionId, ws) => {
    get().wsMap.set(sessionId, ws);
  },
  getWS: (sessionId: string) => {
    return get().wsMap.get(sessionId);
  },
  closeWS: (sessionId) => {
    const ws = get().wsMap.get(sessionId);
    if (ws) {
      ws.close();
      get().wsMap.delete(sessionId);
    }
  },
}));
